'''
# 以医生标注的180个数据作为测试集，便于结果比较
# 提取五分类训练集和测试集
# 提取三分类训练集和测试集
# 提取五分类训练集和测试集

# 未进行数据类别均衡
# 这种数据集划分可能导致模型训练结果出现较大偏差，不符合实际评估
# 因为评估结果是在固定的180个诊断数据上进行，无法保证这180个诊断数据具有代表性
'''


import csv
import pandas as pd
import numpy as np
import shutil
import cv2
import glob
import os

# 提取测试集的病人ID
def readTestId():
    names = pd.read_csv('D:/lung_cancer/data/doctor_result.csv')
    test_id = list(names['patientID'])
    return test_id


# 将所有数据划分为五分类训练集和测试集
def divideData():
    csv_file = open('D:/lung_cancer/data/data_augmentation/all_data5.csv')
    csv_reader_lines = csv.reader(csv_file)
    data = []
    for one_line in csv_reader_lines:
        data.append(one_line)

    test_list = []
    train_list = []

    test_id = readTestId()

    for i in range(1, len(data)):
        if int(data[i][0]) in test_id:
            test_list.append(data[i])
        else:
            train_list.append(data[i])

    print('all data quantity is %d, train is %d, test is %d' %(len(data)-1, len(train_list), len(test_list)))

    save_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/five/'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    df = pd.DataFrame(test_list, columns=data[0])
    df.to_csv(save_path+'test.csv', index=False)

    df = pd.DataFrame(train_list, columns=data[0])
    df.to_csv(save_path+'train.csv', index=False)


# 从五分类测试集中提取出腺癌（1）和鳞癌（2）测试集
def divide_two_test_Data():
    test_csv_file = open('D:/lung_cancer/data/divide_csv/five/test.csv')

    csv_reader_lines = csv.reader(test_csv_file)
    data = []
    for one_line in csv_reader_lines:
        data.append(one_line)
    test_list = []

    # 统计腺癌和鳞癌的测试集数量
    one = 0
    two = 0
    for i in range(1, len(data)):
        if data[i][6] == '1':
            one = one+1
            test_list.append(data[i])
        elif data[i][6] == '2':
            two = two+1
            test_list.append(data[i])

    print('all test data quantity is %d, 1 class test is %d, 2 class test is %d' % (len(data) - 1, one, two))

    df = pd.DataFrame(test_list, columns=data[0])
    df.to_csv('D:/lung_cancer/data/divide_csv/two/test.csv', index=False)


# 从五分类训练集中提取出腺癌（1）和鳞癌（2）训练集
def divide_two_train_Data():
    csv_file = open('D:/lung_cancer/data/divide_csv/five/train.csv')
    csv_reader_lines = csv.reader(csv_file)
    data = []
    for one_line in csv_reader_lines:
        data.append(one_line)
    train_list = []

    # 统计腺癌和鳞癌的训练集数量
    one = 0
    two = 0
    for i in range(1, len(data)):
        if data[i][6] == '1':
            one = one + 1
            train_list.append(data[i])
        elif data[i][6] == '2':
            two = two + 1
            train_list.append(data[i])

    print('all train data quantity is %d, 1 class train is %d, 2 class train is %d' % (len(data) - 1, one, two))

    df = pd.DataFrame(train_list, columns=data[0])
    df.to_csv('D:/lung_cancer/data/divide_csv/two/train.csv', index=False)


# 将五分类训练集变换为三分类训练集（腺癌（1），鳞癌（2）， 其他（3））
def divide_three_train_Data():
    print('Hello')
    csv_file = open('D:/lung_cancer/data/divide_csv/five/train.csv')
    csv_reader_lines = csv.reader(csv_file)
    data = []
    for one_line in csv_reader_lines:
        data.append(one_line)
    train_list = []

    # 统计腺癌和鳞癌和其他的训练集数量
    one = 0
    two = 0
    three = 0
    for i in range(1, len(data)):
        if data[i][6] == '1':
            one = one + 1
            train_list.append(data[i])
        elif data[i][6] == '2':
            two = two + 1
            train_list.append(data[i])
        else:
            three = three+1
            data[i][6] = '3'
            train_list.append(data[i])

    print('all train data quantity is %d, 1 class train is %d, 2 class train is %d, 3 class train is %d' % (len(data) - 1, one, two, three))

    df = pd.DataFrame(train_list, columns=data[0])
    df.to_csv('D:/lung_cancer/data/divide_csv/three/train.csv', index=False)


# 将五分类测试集变换为三分类测试集（腺癌（1），鳞癌（2）， 其他（3））
def divide_three_test_Data():
    print('Hello')
    print('Hello')
    csv_file = open('D:/lung_cancer/data/divide_csv/five/test.csv')
    csv_reader_lines = csv.reader(csv_file)
    data = []
    for one_line in csv_reader_lines:
        data.append(one_line)
    test_list = []

    # 统计腺癌和鳞癌和其他的训练集数量
    one = 0
    two = 0
    three = 0
    for i in range(1, len(data)):
        if data[i][6] == '1':
            one = one + 1
            test_list.append(data[i])
        elif data[i][6] == '2':
            two = two + 1
            test_list.append(data[i])
        else:
            three = three + 1
            data[i][6] = '3'
            test_list.append(data[i])

    print('all test data quantity is %d, 1 class test is %d, 2 class test is %d, 3 class test is %d' % (
    len(data) - 1, one, two, three))

    df = pd.DataFrame(test_list, columns=data[0])
    df.to_csv('D:/lung_cancer/data/divide_csv/three/test.csv', index=False)


# 依照五分类csv文件分离训练集Slice
def divide_five_train_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/five/train.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'

    save_CT_path = 'D:/lung_cancer/data/Slice/five/train/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/five/train/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/five/train/SUVSlice/'

    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)

    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = str(patient)+'_'+str(cancer_type)+'_CTSlice.npy'
        pet_name = str(patient) + '_' + str(cancer_type) + '_PETSlice.npy'
        suv_name = str(patient)+'_'+str(cancer_type)+'_SUVSlice.npy'
        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + ct_name)
        np.save(save_PET_path+pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))


# 依照五分类csv文件分离测试集Slice
def divide_five_test_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/five/test.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'

    save_CT_path = 'D:/lung_cancer/data/Slice/five/test/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/five/test/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/five/test/SUVSlice/'
    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)
    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = str(patient)+'_'+str(cancer_type)+'_CTSlice.npy'
        pet_name = str(patient) + '_' + str(cancer_type) + '_PETSlice.npy'
        suv_name = str(patient) + '_' + str(cancer_type) + '_SUVSlice.npy'
        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + ct_name)
        np.save(save_PET_path+pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))


# 依照二分类csv文件分离训练集Slice
def divide_two_train_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/two/train.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'

    save_CT_path = 'D:/lung_cancer/data/Slice/two/train/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/two/train/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/two/train/SUVSlice/'
    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)
    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = str(patient)+'_'+str(cancer_type)+'_CTSlice.npy'
        pet_name = str(patient) + '_' + str(cancer_type) + '_PETSlice.npy'
        suv_name = str(patient) + '_' + str(cancer_type) + '_SUVSlice.npy'
        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + ct_name)
        np.save(save_PET_path+pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))


# 依照二分类csv文件分离测试集Slice
def divide_two_test_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/two/test.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'

    save_CT_path = 'D:/lung_cancer/data/Slice/two/test/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/two/test/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/two/test/SUVSlice/'
    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)
    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = str(patient)+'_'+str(cancer_type)+'_CTSlice.npy'
        pet_name = str(patient) + '_' + str(cancer_type) + '_PETSlice.npy'
        suv_name = str(patient) + '_' + str(cancer_type) + '_SUVSlice.npy'
        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + ct_name)
        np.save(save_PET_path+pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))


# 依照三分类csv文件分离训练集Slice
def divide_three_train_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/three/train.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'
    save_CT_path = 'D:/lung_cancer/data/Slice/three/train/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/three/train/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/three/train/SUVSlice/'
    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)
    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = glob.glob(origin_CT_slice+str(patient)+'*.npy')[0].split('\\')[-1]
        pet_name = glob.glob(origin_PET_slice+str(patient)+'*.npy')[0].split('\\')[-1]
        suv_name = glob.glob(origin_SUV_slice+str(patient)+'*.npy')[0].split('\\')[-1]
        # print(ct_name)
        # print(pet_name)
        save_ct_name = ct_name
        save_pet_name = pet_name
        save_suv_name = suv_name
        if int(cancer_type) != 1 and int(cancer_type) != 2:
            save_ct_name = str(patient)+'_3_CTSlice.npy'
            save_pet_name = str(patient)+'_3_PETSlice.npy'
            save_suv_name = str(patient)+'_3_SUVSlice.npy'

        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + save_ct_name)
        np.save(save_PET_path+save_pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + save_suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))

# 依照三分类csv文件分离测试集Slice
def divide_three_test_slice():
    src_path = 'D:/lung_cancer/data/divide_csv/three/test.csv'
    origin_CT_slice = 'D:/lung_cancer/data/Slice/CTSlice/'
    origin_PET_slice = 'D:/lung_cancer/data/Slice/PETSlice/'
    origin_SUV_slice = 'D:/lung_cancer/data/Slice/SUVSlice/'

    save_CT_path = 'D:/lung_cancer/data/Slice/three/test/CTSlice/'
    save_PET_path = 'D:/lung_cancer/data/Slice/three/test/PETSlice/'
    save_SUV_path = 'D:/lung_cancer/data/Slice/three/test/SUVSlice/'
    if not os.path.exists(save_CT_path):
        os.makedirs(save_CT_path)
    if not os.path.exists(save_PET_path):
        os.makedirs(save_PET_path)
    if not os.path.exists(save_SUV_path):
        os.makedirs(save_SUV_path)
    data = pd.read_csv(src_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        cancer_type = data['cancer_type'][i]
        ct_name = glob.glob(origin_CT_slice+str(patient)+'*.npy')[0].split('\\')[-1]
        pet_name = glob.glob(origin_PET_slice+str(patient)+'*.npy')[0].split('\\')[-1]
        suv_name = glob.glob(origin_SUV_slice + str(patient) + '*.npy')[0].split('\\')[-1]
        # print(ct_name)
        # print(pet_name)
        save_ct_name = ct_name
        save_pet_name = pet_name
        save_suv_name = suv_name
        if int(cancer_type) != 1 and int(cancer_type) != 2:
            save_ct_name = str(patient)+'_3_CTSlice.npy'
            save_pet_name = str(patient)+'_3_PETSlice.npy'
            save_suv_name = str(patient) + '_3_SUVSlice.npy'
        pet = np.load(origin_PET_slice+pet_name)
        max_pixel = np.max(pet)
        min_pixel = np.min(pet)
        pet = (pet-min_pixel)/float(max_pixel)
        shutil.copyfile(origin_CT_slice + ct_name, save_CT_path + save_ct_name)
        np.save(save_PET_path+save_pet_name, pet)
        shutil.copyfile(origin_SUV_slice + suv_name, save_SUV_path + save_suv_name)
        print('%d--->%s is ok!' % (i, str(patient)))


if __name__ == '__main__':
    print('----------数据切分-----------')
    divideData()
    # divide_two_test_Data()
    # divide_two_train_Data()
    # divide_five_train_slice()
    # divide_five_test_slice()
    # divide_two_train_slice()
    # divide_two_test_slice()
    # divide_three_train_Data()
    # divide_three_test_Data()
    # divide_three_train_slice()
    # divide_three_test_slice()